package com.itcast.stargraph.service.impl;

import cn.hutool.json.JSONUtil;
import com.itcast.stargraph.pojo.ComfyuiTask;
import com.itcast.stargraph.pojo.MessageBase;
import com.itcast.stargraph.service.ComfyuiMsgService;
import com.itcast.stargraph.service.ISgUserResultService;
import com.itcast.stargraph.service.TextToImageService;
import com.itcast.stargraph.service.WsNoticeService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 根据不同情况处理要推送给客户端的消息
 *
 * @Author katana
 * @Date 2025/6/12 15:26
 */
@Service
public class ComfyuiMsgServiceImpl implements ComfyuiMsgService {
    
    @Autowired
    private WsNoticeService wsNoticeService;
    
    @Autowired
    private TextToImageService textToImageService;
    
    @Autowired
    private ISgUserResultService userResultService;
    
    /**
     * 根据不同情况处理要推送给客户端的消息
     *
     * @param msg
     */
    @Override
    public void handleMsg(MessageBase msg) {
        //处理进度消息
        if("progress".equals(msg.getType())) {
            this.handleProgress(msg);
        }else if("executed".equals(msg.getType())) {
            //处理中断消息
            this.handleExecuted(msg);
        }else if("execution_error".equals(msg.getType())) {
            //处理中断消息
            this.handleExecutedError(msg);
        }
    }
    
    private void handleExecutedError(MessageBase msg) {
        //1. 获取消息里的任务id
        HashMap<String, Object> data = msg.getData();
        Object promptId = data.getOrDefault("prompt_id", null);
        if(promptId==null) {
            return;
        }
        
        //2. 获取缓存中最开始的任务
        ComfyuiTask comfyuiTask = textToImageService.getStartedTask(promptId.toString());
        if(comfyuiTask==null){
            return;
        }
        
        //3. 组装前端需要的消息
        data.put("type","execution_error");
        
        //4. 推送消息发给前端
        wsNoticeService.sendToUser(comfyuiTask.getWsClientId(), JSONUtil.toJsonStr(data));
    }
    
    private void handleExecuted(MessageBase msg) {
        //1. 获取消息里的任务id
        HashMap<String, Object> data = msg.getData();
        Object promptId = data.getOrDefault("prompt_id", null);
        if(promptId==null) {
            return;
        }
        
        //2. 获取缓存中最开始的任务
        ComfyuiTask comfyuiTask = textToImageService.getStartedTask(promptId.toString());
        if(comfyuiTask==null){
            return;
        }
        //3 组装前端需要的消息
        data.put("type","imageResult");
        Object outputObj = data.get("output");
        if (outputObj==null){
            return;
        }
        Map<String, Object> output = (Map<String, Object>)outputObj;
        List<Map<String,Object>> images = (List<Map<String,Object>> )output.get("images");
        List<String> imageUrls = images.stream().map(image -> {
            return String.format("http://192.168.100.1:8188/view?filename=%s&type=%s&&subfolder="
                    , image.get("filename"), image.get("type"));
        }).toList();
        data.put("urls" ,imageUrls);
        
        
        //4. 保存用户图片
        userResultService.insertUserResult(comfyuiTask.getUserId(), imageUrls);
        //5. 推送消息发给前端
        wsNoticeService.sendToUser(comfyuiTask.getWsClientId(), JSONUtil.toJsonStr(data));
    }
    
    
    private void handleProgress(MessageBase msg) {
        //获取进度消息中的promptId
        HashMap<String, Object> data = msg.getData();
        String promptId = (String) data.get("prompt_id");
        if(promptId==null) {
            return;
        }
        //获取进度消息中的进度值
        //2. 获取缓存中优先级最高的任务
        ComfyuiTask comfyuiTask = textToImageService.getStartedTask(promptId.toString());
        if(comfyuiTask==null){
            return;
        }
        //3. 组装前端需要的消息
        data.put("type","progress");
        //4. 推送消息发给前端
        wsNoticeService.sendToUser(comfyuiTask.getWsClientId(), JSONUtil.toJsonStr(data));
    }
}
